"""
@Time: 2020/12/1 下午 8:15
@Author: jinzhuan
@File: trex_ner_test.py
@Desc: 
"""
import sys
sys.path.append('/data/zhuoran/code/cognlp')
sys.path.append('/data/zhuoran/cognlp')

from cognlp import *
from cognlp.core.dataset import TrexNerDataset
from cognlp.io.processor.ner.conll2003 import TrexNerProcessor
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import RandomSampler
from cognlp.core.trainer import Tester

if __name__ == '__main__':
    # loader = TrexNerLoader()
    # data = loader.load_all(path='../data/ner/trex/data')
    # processor = TrexNerProcessor(loader.get_labels(), path='../data/ner/trex/data')
    # train_datable = processor.process(data)
    # train_datable.print_table(5)
    # train_datable.save_table("../data/ner/trex/data/train.json")
    torch.cuda.set_device(4)
    device = torch.device('cuda')
    processor = TrexNerProcessor(path='../data/ner/trex/data')
    train_datable = DataTable.load_table("../data/ner/trex/data/train.json")
    train_data = TrexNerDataset(train_datable, device=device)
    train_sampler = RandomSampler(train_data)
    model = Bert4Ner(len(processor.vocabulary), device=device)
    loss = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.000001)
    metric = SpanFPreRecMetric(tag_vocab=processor.vocabulary)

    tester = Tester(model, model_path='../data/ner/trex/model/2020-12-04-03:46:31-model.pkl', batch_size=16, sampler=train_sampler,
                     drop_last=False, num_workers=0, print_every=1000,
                     dev_data=train_data, metrics=metric, metric_key=None, use_tqdm=True, device=None,
                     callbacks=None, check_code_level=0, device_ids=[4, 5, 6])
    tester.test()
    # trainer = Trainer(train_data=train_data, model=model, optimizer=optimizer, loss=loss,
    #                   batch_size=64, train_sampler=train_sampler, drop_last=False, gradient_accumulation_steps=1,
    #                   num_workers=5, n_epochs=10, print_every=50, scheduler=None, dev_sampler=train_sampler,
    #                   dev_data=train_data, metrics=metric, metric_key=None, validate_steps=1,
    #                   save_path="../data/ner/trex/model", save_file=None, save_steps=1, use_tqdm=True,
    #                   device=device,
    #                   callbacks=None, check_code_level=0, grad_norm=None, device_ids=[4, 5, 6])
    # trainer.train()

